#!/bin/bash

# CPU工作负载模拟器
# 用法: ./cpu_workload_simulator.sh [工作负载类型] [持续时间(秒)]

# 默认参数
workload_type=${1:-"mixed"}  # 默认混合负载
duration=${2:-60}           # 默认运行60秒

# 检查必要的命令是否存在
check_command() {
    command -v "$1" &> /dev/null || { echo "错误: 需要 $1 命令但未找到"; exit 1; }
}

check_command bc

# 输出配置信息
echo "CPU工作负载模拟器"
echo "负载类型: $workload_type"
echo "持续时间: $duration 秒"
echo "按 Ctrl+C 可随时终止测试"

# 记录开始时间
start_time=$(date +%s)
end_time=$((start_time + duration))

# 定义各种工作负载函数

# 计算密集型负载 - 大量数学计算
compute_intensive() {
    echo "执行计算密集型负载..."
    while [ $(date +%s) -lt $end_time ]
    do
        # 执行复杂的数学计算
        echo "scale=50; 4*a(1)" | bc -l > /dev/null
        echo "scale=50; s(1)*c(1)" | bc -l > /dev/null
    done
}

# I/O密集型负载 - 读写文件
io_intensive() {
    echo "执行I/O密集型负载..."
    tmp_file="/tmp/io_test_$$.dat"
    while [ $(date +%s) -lt $end_time ]
    do
        # 写入和读取文件
        dd if=/dev/urandom of=$tmp_file bs=1M count=10 2>/dev/null
        cat $tmp_file > /dev/null
        rm $tmp_file
    done
}

# 内存密集型负载 - 大量内存分配和释放
memory_intensive() {
    echo "执行内存密集型负载..."
    while [ $(date +%s) -lt $end_time ]
    do
        # 生成大字符串并存储在变量中
        local big_data=$(cat /dev/urandom | tr -dc 'a-zA-Z0-9' | head -c 100000)
        # 使用变量防止被优化掉
        echo ${#big_data} > /dev/null
        unset big_data
    done
}

# 混合型负载 - 组合上述负载类型
mixed_workload() {
    echo "执行混合负载类型..."
    while [ $(date +%s) -lt $end_time ]
    do
        # 计算
        echo "scale=20; 4*a(1)" | bc -l > /dev/null
        
        # I/O
        tmp_file="/tmp/mixed_test_$$.dat"
        dd if=/dev/urandom of=$tmp_file bs=512k count=1 2>/dev/null
        cat $tmp_file > /dev/null
        rm $tmp_file
        
        # 内存
        local data=$(cat /dev/urandom | tr -dc 'a-zA-Z0-9' | head -c 50000)
        echo ${#data} > /dev/null
        unset data
    done
}

# 突发型负载 - 短时高负载，然后休眠
bursty_workload() {
    echo "执行突发型负载..."
    while [ $(date +%s) -lt $end_time ]
    do
        # 高负载阶段
        for i in {1..5}; do
            echo "scale=30; 4*a(1)" | bc -l > /dev/null
        done
        
        # 休眠阶段
        sleep 2
    done
}

# 根据指定的工作负载类型执行相应的函数
case "$workload_type" in
    "compute")
        compute_intensive
        ;;
    "io")
        io_intensive
        ;;
    "memory")
        memory_intensive
        ;;
    "bursty")
        bursty_workload
        ;;
    "mixed"|*)
        mixed_workload
        ;;
esac

echo "CPU工作负载模拟完成!" 